import torch
import mindspore as ms


# 通过PyTorch参数文件，打印PyTorch的参数文件里所有参数的参数名和shape，返回参数字典
def pytorch_params(pth_file):
    par_dict = torch.load(pth_file, map_location='cpu')
    pt_params = {}
    for name in par_dict:
        parameter = par_dict[name]
        print(name, parameter.numpy().shape)
        pt_params[name] = parameter.numpy()
    return pt_params

# 通过MindSpore的Cell，打印Cell里所有参数的参数名和shape，返回参数字典
def mindspore_params(network):
    ms_params = {}
    for param in network.values():
        name = param.name
        value = param.data.asnumpy()
        print(name, value.shape)
        ms_params[name] = value
    return ms_params

if __name__ == "__main__":
    print("="*20+"Generator"+"="*20)
    pth_path = "/root/linx/segan_pt/pt_generator.pth"
    pt_param = pytorch_params(pth_path)
    print("="*20)
    ms_param = mindspore_params(ms.load_checkpoint("/root/linx/segan_ms/converted_g.ckpt"))

    # print("\n"+"="*20+"Discriminator"+"="*20)
    # pth_path = "/root/linx/segan_pt/pt_disciminator.pth"
    # pt_param = pytorch_params(pth_path)
    # print("="*20)
    # ms_param = mindspore_params(ms.load_checkpoint("/root/linx/segan_ms/ms_disciminator.ckpt"))

    # print("="*20+"Generator"+"="*20)
    # ms_param = mindspore_params(ms.load_checkpoint("/root/linx/segan_ms/converted_d.ckpt"))
    # print("="*20)
    # ms_param = mindspore_params(ms.load_checkpoint("/root/linx/segan_ms/ms_disciminator.ckpt"))